package xyz.labmem.lipx.client.core.handler

import cn.hutool.core.codec.Base64
import cn.hutool.core.convert.Convert
import cn.hutool.crypto.Mode
import cn.hutool.crypto.Padding
import cn.hutool.crypto.symmetric.AES
import io.netty.channel.Channel
import io.netty.channel.ChannelHandlerContext
import io.netty.channel.ChannelInitializer
import io.netty.channel.group.DefaultChannelGroup
import io.netty.channel.socket.SocketChannel
import io.netty.handler.codec.bytes.ByteArrayDecoder
import io.netty.handler.codec.bytes.ByteArrayEncoder
import io.netty.util.concurrent.GlobalEventExecutor
import xyz.labmem.lipx.core.context.ServerContext
import xyz.labmem.lipx.core.entity.ClientConfig
import xyz.labmem.lipx.core.handler.LabCommonHandler
import xyz.labmem.lipx.core.handler.impl.LocalProxyHandler
import xyz.labmem.lipx.core.handler.impl.RemoteProxyHandler
import xyz.labmem.lipx.core.net.TcpConnection
import xyz.labmem.lipx.core.net.TcpServer
import xyz.labmem.lipx.core.protocol.LabMessage
import xyz.labmem.lipx.core.protocol.LabMessageType
import xyz.labmem.lipx.core.tool.LogInfo
import xyz.labmem.lipx.core.tool.pwdJoinSalt
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Executors
import java.util.concurrent.TimeUnit

/**
 * @description: do something
 * @author: ltty
 * @date: 2022/5/30 15:58
 */
class LabNettyClientHandler(
    private val config: ClientConfig,
    private val isReceive: Boolean,//是否正向代理
    private val registerBack: ((r: Boolean, reject: Boolean, remake: Boolean, reason: String) -> Unit)
) : LabCommonHandler() {

    private var register = false
    private val channelHandlerMap = ConcurrentHashMap<String, LabCommonHandler>()
    private val channelGroup = DefaultChannelGroup(GlobalEventExecutor.INSTANCE)
    private var aes: AES? = null
    private var remoteConnectionServer: TcpServer? = null

    @Throws(Exception::class)
    override fun channelActive(ctx: ChannelHandlerContext) {
        val metaData = HashMap<String, Any>()
        metaData["id"] = config.id!!
        if (isReceive) {
            metaData["proxyHost"] = config.proxyHost!!
            config.proxyPort?.let {
                metaData["proxyPort"] = it
            }
        } else {
            config.targetPort?.let {
                metaData["port"] = it
            }
        }
        metaData["password"] = pwdJoinSalt(config.password!!)
        metaData["wls"] = config.wls
        metaData["isReceive"] = isReceive
        metaData["name"] = config.remark ?: "NoName"
        // register client information
        LabMessage().apply {
            type = LabMessageType.REGISTER
            this.metaData = metaData
            ctx.writeAndFlush(this)
            super.channelActive(ctx)
        }
    }

    @Throws(Exception::class)
    override fun channelRead(ctx: ChannelHandlerContext, msg: Any) {
        val message = msg as LabMessage
        if (message.type === LabMessageType.REGISTER_RESULT) {
            processRegisterResult(message)
        } else if (message.type === LabMessageType.CONNECTED) {
            processConnected(message)
        } else if (message.type === LabMessageType.DISCONNECTED) {
            processDisconnected(message)
        } else if (message.type === LabMessageType.DATA) {
            processData(message)
        } else if (message.type === LabMessageType.KEEPALIVE) {
            // 心跳包, 不处理
        } else {
            throw Exception("Unknown type: ${message.type}")
        }
    }

    @Throws(Exception::class)
    override fun channelInactive(ctx: ChannelHandlerContext) {
        if (isReceive) {
            remoteConnectionServer?.close()
        } else
            channelGroup.close()
        if (!register)
            registerBack(false, false, true, "与服务器已断开连接")
        LogInfo.appendLog("与服务器已断开连接!", config)
    }

    private fun processRegisterResult(message: LabMessage) {
        var reject = false
        if (message.metaData!!["success"] as Boolean) {
            register = true
            config.version = message.metaData!!["version"] as String
            if (message.metaData!!.containsKey("SecretKey") && message.metaData!!.containsKey("vi"))
                aes = AES(
                    Mode.CTR,
                    Padding.NoPadding,
                    Base64.decode((message.metaData!!["SecretKey"]!! as String)),
                    Base64.decode((message.metaData!!["vi"]!! as String))
                )
            LogInfo.appendLog("客户端连接服务器成功", config)
            if (isReceive) {
                try {
                    remoteConnectionServer = TcpServer()
                    remoteConnectionServer!!.bind(
                        config.targetPort!!,
                        object : ChannelInitializer<SocketChannel>() {
                            @Throws(Exception::class)
                            override fun initChannel(ch: SocketChannel) {
                                ch.pipeline().apply {
                                    addLast(
                                        ByteArrayDecoder(),
                                        ByteArrayEncoder(),
                                        RemoteProxyHandler(this@LabNettyClientHandler, config.targetPort!!, aes!!)
                                    )
                                }
                                ServerContext.channels.add(ch)
                            }
                        })
                } catch (e: Exception) {
                    register = false
                    LogInfo.appendLogError("客户端启动本地端口监听失败：${e.message}", config)
                }
            }
        } else {
            register = false
            if (message.metaData!!["reason"].toString().contains("拒绝连接"))
                reject = true
            LogInfo.appendLogError("客户端连接失败：${message.metaData!!["reason"]}", config)
            ctx!!.close()
        }
        if (register) {
            Executors.newSingleThreadScheduledExecutor().apply {
                scheduleAtFixedRate({
                    LabMessage().apply {
                        type = LabMessageType.KEEPALIVE
                        ctx!!.writeAndFlush(this)
                    }
                }, 0, 2000, TimeUnit.MILLISECONDS)
            }
        }
        registerBack(register, reject, false, message.metaData!!["reason"].toString())
    }


    @Throws(Exception::class)
    private fun processConnected(message: LabMessage) {
        try {
            val thisHandler = this
            val localConnection = TcpConnection()
            localConnection.connect(
                config.proxyHost!!,
                Convert.toInt(config.proxyPort!!),
                object : ChannelInitializer<SocketChannel>() {
                    @Throws(Exception::class)
                    override fun initChannel(ch: SocketChannel) {
                        val localProxyHandler =
                            LocalProxyHandler(thisHandler, message.metaData!!["channelId"].toString(), aes!!)
                        ch.pipeline().addLast(ByteArrayDecoder(), ByteArrayEncoder(), localProxyHandler)
                        channelHandlerMap[message.metaData!!["channelId"].toString()] = localProxyHandler
                        channelGroup.add(ch)
                    }
                })
        } catch (e: Exception) {
            val metaData = HashMap<String, Any>()
            metaData["channelId"] = message.metaData!!["channelId"].toString()
            LabMessage().apply {
                type = LabMessageType.DISCONNECTED
                this.metaData = metaData
                ctx!!.writeAndFlush(this)
            }
            channelHandlerMap.remove(message.metaData!!["channelId"])
            e.printStackTrace()
            throw Exception("连接错误断开连接")
        }
    }

    /**
     * DISCONNECTED
     */
    private fun processDisconnected(message: LabMessage) {
        if (isReceive) {
            ServerContext.channels.close { channel: Channel ->
                channel.id().asLongText() == message.metaData!!["channelId"]
            }
        } else {
            val channelId: String = message.metaData!!["channelId"].toString()
            val handler = channelHandlerMap[channelId]
            if (handler != null) {
                handler.ctx?.close()
                channelHandlerMap.remove(channelId)
            }
        }
    }

    /**
     * DATA
     */
    private fun processData(message: LabMessage) {
        if (isReceive) {
            ServerContext.channels.writeAndFlush(
                aes!!.decrypt(message.data)
            ) { channel: Channel ->
                channel.id().asLongText() == message.metaData!!["channelId"]
            }
        } else {
            val channelId: String = message.metaData!!["channelId"].toString()
            val handler = channelHandlerMap[channelId]
            if (handler != null) {
                val ctx = handler.ctx
                ctx?.writeAndFlush(aes?.decrypt(message.data) ?: message.data)
            }
        }
    }
}